# encodign=utf-8
import time
import cv2
import onnxruntime
import numpy

def array2bytes(array_img, suffix):
    # coding img
    success, encoded_array = cv2.imencode("." + suffix, array_img)
    # to bytes
    bytes_img = encoded_array.tobytes()

    return bytes_img

t1 = time.perf_counter()
session = onnxruntime.InferenceSession("export_funds_stillgan.onnx")
t2 = time.perf_counter()


ori_img = cv2.imread("test.jpg")
# preprocess
ori_size = ori_img.shape[:2]

trans_image = cv2.cvtColor(ori_img, cv2.COLOR_BGR2RGB)
trans_image = cv2.resize(trans_image, (512, 512))
trans_image = trans_image / 255.0
img_arr = (trans_image - 0.5) / 0.5

x = numpy.expand_dims(img_arr.astype(numpy.float32).transpose((2, 0, 1)), axis=0)
# infer

inputs = {"input":  x}
outputs = session.run(None, inputs)
print(outputs[0].shape)

H, W = ori_size
image = (numpy.transpose(outputs[0][0], (1, 2, 0)) + 1) / 2.0 * 255.0
(r, g, b) = cv2.split(image)
img_arr = cv2.merge([b, g, r])
img_arr = cv2.resize(img_arr, (W, H), interpolation=cv2.INTER_CUBIC)
img_bytes = array2bytes(img_arr, "png")
image = numpy.asarray(bytearray(img_bytes), dtype=numpy.uint8)
img_arr = cv2.imdecode(image, cv2.IMREAD_COLOR)

cv2.imwrite("out.png", img_arr)
t3 = time.perf_counter()
print(f"init cost:{t2-t1}\ninfer cost:{t3-t2}")